{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Number of Good Paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #union-find #graph #array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #并查集 #图 #数组"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: numberOfGoodPaths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #好路径的数目"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一棵 <code>n</code>&nbsp;个节点的树（连通无向无环的图），节点编号从&nbsp;<code>0</code>&nbsp;到&nbsp;<code>n - 1</code>&nbsp;且恰好有&nbsp;<code>n - 1</code>&nbsp;条边。</p>\n",
    "\n",
    "<p>给你一个长度为 <code>n</code>&nbsp;下标从 <strong>0</strong>&nbsp;开始的整数数组&nbsp;<code>vals</code>&nbsp;，分别表示每个节点的值。同时给你一个二维整数数组&nbsp;<code>edges</code>&nbsp;，其中&nbsp;<code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code>&nbsp;表示节点&nbsp;<code>a<sub>i</sub></code> 和&nbsp;<code>b<sub>i</sub></code><sub>&nbsp;</sub>之间有一条&nbsp;<strong>无向</strong>&nbsp;边。</p>\n",
    "\n",
    "<p>一条 <strong>好路径</strong>&nbsp;需要满足以下条件：</p>\n",
    "\n",
    "<ol>\n",
    "\t<li>开始节点和结束节点的值 <strong>相同</strong>&nbsp;。</li>\n",
    "\t<li>开始节点和结束节点中间的所有节点值都 <strong>小于等于</strong>&nbsp;开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。</li>\n",
    "</ol>\n",
    "\n",
    "<p>请你返回不同好路径的数目。</p>\n",
    "\n",
    "<p>注意，一条路径和它反向的路径算作 <strong>同一</strong>&nbsp;路径。比方说，&nbsp;<code>0 -&gt; 1</code>&nbsp;与&nbsp;<code>1 -&gt; 0</code>&nbsp;视为同一条路径。单个节点也视为一条合法路径。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例&nbsp;1：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2022/08/04/f9caaac15b383af9115c5586779dec5.png\" style=\"width: 400px; height: 333px;\"></p>\n",
    "\n",
    "<pre><b>输入：</b>vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]\n",
    "<b>输出：</b>6\n",
    "<b>解释：</b>总共有 5 条单个节点的好路径。\n",
    "还有 1 条好路径：1 -&gt; 0 -&gt; 2 -&gt; 4 。\n",
    "（反方向的路径 4 -&gt; 2 -&gt; 0 -&gt; 1 视为跟 1 -&gt; 0 -&gt; 2 -&gt; 4 一样的路径）\n",
    "注意 0 -&gt; 2 -&gt; 3 不是一条好路径，因为 vals[2] &gt; vals[0] 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2022/08/04/149d3065ec165a71a1b9aec890776ff.png\" style=\"width: 273px; height: 350px;\"></p>\n",
    "\n",
    "<pre><b>输入：</b>vals = [1,1,2,2,3], edges = [[0,1],[1,2],[2,3],[2,4]]\n",
    "<b>输出：</b>7\n",
    "<strong>解释：</strong>总共有 5 条单个节点的好路径。\n",
    "还有 2 条好路径：0 -&gt; 1 和 2 -&gt; 3 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 3：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2022/08/04/31705e22af3d9c0a557459bc7d1b62d.png\" style=\"width: 100px; height: 88px;\"></p>\n",
    "\n",
    "<pre><b>输入：</b>vals = [1], edges = []\n",
    "<b>输出：</b>1\n",
    "<b>解释：</b>这棵树只有一个节点，所以只有一条好路径。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>n == vals.length</code></li>\n",
    "\t<li><code>1 &lt;= n &lt;= 3 * 10<sup>4</sup></code></li>\n",
    "\t<li><code>0 &lt;= vals[i] &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>edges.length == n - 1</code></li>\n",
    "\t<li><code>edges[i].length == 2</code></li>\n",
    "\t<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub> &lt; n</code></li>\n",
    "\t<li><code>a<sub>i</sub> != b<sub>i</sub></code></li>\n",
    "\t<li><code>edges</code>&nbsp;表示一棵合法的树。</li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [number-of-good-paths](https://leetcode.cn/problems/number-of-good-paths/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [number-of-good-paths](https://leetcode.cn/problems/number-of-good-paths/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[1,3,2,1,3]\\n[[0,1],[0,2],[2,3],[2,4]]', '[1,1,2,2,3]\\n[[0,1],[1,2],[2,3],[2,4]]', '[1]\\n[]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "     \n",
    "        n = len(vals)\n",
    "        record = [[] for i in range(n)]\n",
    "        for ele in edges:\n",
    "            record[ele[0]].append(ele[1])\n",
    "            record[ele[1]].append(ele[0])\n",
    "\n",
    "        fa = [i for i in range(n)]   # 根节点存最大的节点\n",
    "        def find_root(x:int)->int:\n",
    "            root = x\n",
    "            while fa[root]!=root:\n",
    "                root = fa[root]\n",
    "            fa[x] = root\n",
    "            return root\n",
    "\n",
    "        size = [1]*n\n",
    "        output = n\n",
    "        for value, ops in sorted(zip(vals, range(n))):\n",
    "            f_ops = find_root(ops)\n",
    "            for neighbor in record[ops]:\n",
    "                f_neighbor = find_root(neighbor)\n",
    "\n",
    "                if f_neighbor==f_ops or vals[f_neighbor]>value:\n",
    "                    continue\n",
    "\n",
    "                if vals[f_neighbor]==value:\n",
    "                    output += size[f_ops] * size[f_neighbor]\n",
    "                    size[f_ops] += size[f_neighbor]\n",
    "                \n",
    "\n",
    "                fa[f_neighbor] = f_ops\n",
    "\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        nexts = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            nexts[a].append(b)\n",
    "            nexts[b].append(a)\n",
    "        parents = list(range(n))\n",
    "        size = [1]*n\n",
    "        res = n\n",
    "        def parent(i):\n",
    "            stack = []\n",
    "            while parents[i] != i:\n",
    "                stack.append(i)\n",
    "                i = parents[i]\n",
    "            for j in stack:\n",
    "                parents[j] = i\n",
    "            return i\n",
    "        for v, i in sorted(zip(vals, range(n))):\n",
    "            pi = parent(i)\n",
    "            for nex in nexts[i]:\n",
    "                nex = parent(nex)\n",
    "                if nex == pi or vals[nex] > v: continue\n",
    "                if vals[nex] == v:\n",
    "                    res += size[nex] * size[pi]\n",
    "                    size[pi] += size[nex]\n",
    "                parents[nex] = pi\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    @ staticmethod\n",
    "    def myFind(parent, i):\n",
    "        if parent[i] != i:\n",
    "            parent[i] = Solution.myFind(parent, parent[i])\n",
    "        return parent[i]\n",
    "\n",
    "    @ staticmethod\n",
    "    def myUnion(parent, rank, i, j):\n",
    "        pi, pj = Solution.myFind(parent, i), Solution.myFind(parent, j)\n",
    "        if rank[pi] > rank[pj]:\n",
    "            parent[pj] = pi\n",
    "        elif rank[pi] < rank[pj]:\n",
    "            parent[pi] = pj\n",
    "        else:\n",
    "            parent[pj] = pi\n",
    "            rank[pi] += 1\n",
    "        return\n",
    "\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        max_edge = {}\n",
    "        for i, j in edges:\n",
    "            m = max(vals[i], vals[j])\n",
    "            if m not in max_edge: \n",
    "                max_edge[m] = []\n",
    "            max_edge[m].append([i, j])\n",
    "        parent, rank = [_ for _ in range(n)], [1] * n\n",
    "        ans = 0\n",
    "        for key in sorted(list(max_edge.keys())):\n",
    "            s = set()\n",
    "            for i, j in max_edge[key]:\n",
    "                if vals[i] == key: s.add(i)\n",
    "                if vals[j] == key: s.add(j)\n",
    "                Solution.myUnion(parent, rank, i, j)\n",
    "            tmp = []\n",
    "            for i in s:\n",
    "                tmp.append(Solution.myFind(parent, i))\n",
    "            c = Counter(tmp)\n",
    "            for k in c:\n",
    "                ans += (c[k] * (c[k] - 1)) // 2\n",
    "        return ans + n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        parent = {i : i for i in range(n)}\n",
    "        def find(node):\n",
    "            path = []\n",
    "            while node != parent[node]:\n",
    "                path.append(node)\n",
    "                node = parent[node]\n",
    "            for p in path:\n",
    "                parent[p] = node \n",
    "            return node \n",
    "        \n",
    "        def union(n1, n2):\n",
    "            p1 = find(n1)\n",
    "            p2 = find(n2)\n",
    "            if p1 != p2:\n",
    "                parent[p1] = p2 \n",
    "        \n",
    "        dic = collections.defaultdict(list)\n",
    "        for i in range(n):\n",
    "            dic[vals[i]].append(i)\n",
    "        res = 0\n",
    "        edges.sort(key = lambda x: - max(vals[x[0]], vals[x[1]]))\n",
    "        for val in sorted(dic.keys()):\n",
    "            while edges and max(vals[edges[-1][0]], vals[edges[-1][1]]) == val:\n",
    "                a, b = edges.pop()\n",
    "                union(a, b)\n",
    "            counter = collections.Counter()\n",
    "            for node in dic[val]:\n",
    "                counter[find(node)] += 1\n",
    "            for c in counter:\n",
    "                res += (counter[c] - 1) * counter[c] // 2 \n",
    "        return res + n \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        dic = collections.defaultdict(list)\n",
    "        for i in range(n):\n",
    "            dic[vals[i]].append(i)\n",
    "        parent = {i : i for i in range(n)}\n",
    "        def find(node):\n",
    "            path = []\n",
    "            while node != parent[node]:\n",
    "                path.append(node)\n",
    "                node = parent[node]\n",
    "            for p in path:\n",
    "                parent[p] = node \n",
    "            return node \n",
    "        \n",
    "        def union(n1, n2):\n",
    "            p1 = find(n1)\n",
    "            p2 = find(n2)\n",
    "            if p1 != p2:\n",
    "                parent[p1] = p2 \n",
    "        res = 0\n",
    "        edges.sort(key = lambda x: - max(vals[x[0]], vals[x[1]]))\n",
    "        while edges:\n",
    "            curMax = max(vals[edges[-1][0]], vals[edges[-1][1]])\n",
    "            while edges and max(vals[edges[-1][0]], vals[edges[-1][1]]) == curMax:\n",
    "                a, b = edges.pop()\n",
    "                union(a, b)\n",
    "            counter = collections.Counter()\n",
    "            for node in dic[curMax]:\n",
    "                counter[find(node)] += 1 \n",
    "            for c in counter:\n",
    "                res += counter[c] * (counter[c] - 1) // 2 \n",
    "        return res + n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self.root = [i for i in range(n)]\n",
    "        self.size = [1] * n\n",
    "        self.part = n\n",
    "        return\n",
    "\n",
    "    def find(self, x):\n",
    "        lst = []\n",
    "        while x != self.root[x]:\n",
    "            lst.append(x)\n",
    "            # 在查询的时候合并到顺带直接根节点\n",
    "            x = self.root[x]\n",
    "        for w in lst:\n",
    "            self.root[w] = x\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        root_x = self.find(x)\n",
    "        root_y = self.find(y)\n",
    "        if root_x == root_y:\n",
    "            return False\n",
    "        if self.size[root_x] >= self.size[root_y]:\n",
    "            root_x, root_y = root_y, root_x\n",
    "        self.root[root_x] = root_y\n",
    "        self.size[root_y] += self.size[root_x]\n",
    "        # 将非根节点的秩赋0\n",
    "        self.size[root_x] = 0\n",
    "        self.part -= 1\n",
    "        return True\n",
    "\n",
    "    def is_connected(self, x, y):\n",
    "        return self.find(x) == self.find(y)\n",
    "\n",
    "    def get_root_part(self):\n",
    "        # 获取每个根节点对应的组\n",
    "        part = defaultdict(list)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            part[self.find(i)].append(i)\n",
    "        return part\n",
    "\n",
    "    def get_root_size(self):\n",
    "        # 获取每个根节点对应的组大小\n",
    "        size = defaultdict(int)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            size[self.find(i)] = self.size[self.find(i)]\n",
    "        return size\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        \n",
    "        n = len(vals)\n",
    "        ans = 0\n",
    "        uf = UnionFind(n)\n",
    "        dct = defaultdict(list)\n",
    "        for i in range(n):\n",
    "            dct[vals[i]].append(i)\n",
    "            \n",
    "        edge = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            edge[x].append(y)\n",
    "            edge[y].append(x)\n",
    "            \n",
    "        for num in sorted(dct):\n",
    "            for i in dct[num]:\n",
    "                for j in edge[i]:\n",
    "                    if vals[j] <= num:\n",
    "                        uf.union(i, j)\n",
    "            roots = [uf.find(i) for i in dct[num]]\n",
    "            cnt = Counter(roots)\n",
    "            for va in cnt.values():\n",
    "                ans += va*(va-1)//2\n",
    "            \n",
    "        return ans + n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# sort vertex according to their value, adding edges one by one\r\n",
    "# any vertex that is currently connected with other vertex of same value, forms a good path\r\n",
    "pa = [0] * 30000\r\n",
    "cnt: dict[int, int] = dict()\r\n",
    "\r\n",
    "def find(i: int) -> int:\r\n",
    "    if i != pa[i]:\r\n",
    "        pa[i] = find(pa[i])\r\n",
    "    return pa[i]\r\n",
    "\r\n",
    "def join(i: int, j: int) -> int:\r\n",
    "    f, g = find(i), find(j)\r\n",
    "    if f == g: return 0\r\n",
    "    pa[g] = f\r\n",
    "    ret = cnt[f] * cnt[g] if f in cnt and g in cnt else 0\r\n",
    "    if g in cnt:\r\n",
    "        cnt.setdefault(f, 0)\r\n",
    "        cnt[f] += cnt[g]\r\n",
    "    return ret\r\n",
    "\r\n",
    "class Solution:\r\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\r\n",
    "        n = len(vals)\r\n",
    "        uniq = sorted(set(vals))\r\n",
    "        edges.sort(key=lambda e: max(vals[e[0]], vals[e[1]]))\r\n",
    "        for i in range(n):\r\n",
    "            pa[i] = i\r\n",
    "\r\n",
    "        res = n\r\n",
    "        it = iter(edges)\r\n",
    "        u, v = next(it, [None, None])\r\n",
    "\r\n",
    "        for val in uniq:\r\n",
    "            cnt.clear()\r\n",
    "            while u is not None and max(vals[u], vals[v]) <= val:\r\n",
    "                if vals[u] == val and u not in cnt:\r\n",
    "                    cnt[u] = 1\r\n",
    "                if vals[v] == val and v not in cnt:\r\n",
    "                    cnt[v] = 1\r\n",
    "\r\n",
    "                res += join(u, v)\r\n",
    "                u, v = next(it, [None, None])\r\n",
    "        \r\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:#并查集模板\n",
    "    def __init__(self, n):\n",
    "        self.p = list(range(n))\n",
    "        self.size = [1] * n\n",
    "\n",
    "    def find(self, x):\n",
    "        if self.p[x] != x:\n",
    "            self.p[x] = self.find(self.p[x])\n",
    "        return self.p[x]\n",
    "\n",
    "    def union(self, a, b):\n",
    "        pa, pb = self.find(a), self.find(b)\n",
    "        if pa != pb:\n",
    "            if self.size[pa] >= self.size[pb]:\n",
    "                self.size[pa] += self.size[pb]\n",
    "                self.p[pb] = pa\n",
    "            else:\n",
    "                self.size[pb] += self.size[pa]\n",
    "                self.p[pa] = pb\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        d = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            d[max(vals[u], vals[v])].append([u, v])\n",
    "        uf = UnionFind(n)\n",
    "        ans = n\n",
    "        groups = defaultdict(list)\n",
    "        for i, j in enumerate(vals):\n",
    "            groups[j].append(i)\n",
    "        for i in sorted(d.keys()):\n",
    "            for u, v in d[i]:\n",
    "                uf.union(u, v)\n",
    "            cnt = Counter()\n",
    "            for j in groups[i]:\n",
    "                cnt[uf.find(j)] += 1\n",
    "            for kk, vv in cnt.items():\n",
    "                ans += comb(vv, 2)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, x) -> None:\n",
    "        self.uf = [-1] * x\n",
    "\n",
    "    def find(self, x):\n",
    "        r = x\n",
    "        while self.uf[x] >= 0:\n",
    "            x = self.uf[x]\n",
    "\n",
    "        while r != x:\n",
    "            self.uf[r], r = x, self.uf[r]\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        ux, uy = self.find(x), self.find(y)\n",
    "        if ux == uy:\n",
    "            return\n",
    "        if self.uf[ux] >= self.uf[uy]:\n",
    "            self.uf[uy] += self.uf[ux]\n",
    "            self.uf[ux] = uy\n",
    "        else:\n",
    "            self.uf[ux] += self.uf[uy]\n",
    "            self.uf[uy] = ux\n",
    "        return\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        uf = UnionFind(n)\n",
    "        ans = n\n",
    "        a = defaultdict(list)\n",
    "        for i in range(n):\n",
    "            a[vals[i]].append(i)\n",
    "        b = list(a.items())\n",
    "        b.sort(key = lambda x:x[0])\n",
    "        e = defaultdict(list)\n",
    "        for i in range(n - 1):\n",
    "            u,v = edges[i]\n",
    "            e[u].append([v,max(vals[u],vals[v])])\n",
    "            e[v].append([u,max(vals[v],vals[u])])\n",
    "        for i in range(n):\n",
    "            e[i].sort(key = lambda x:-x[1])\n",
    "        # 从小到大枚举\n",
    "        for v in b:\n",
    "            curv = v[0]\n",
    "            tmp = v[1]\n",
    "            vv = set()\n",
    "            cc = defaultdict(int)\n",
    "            # 合法边加入\n",
    "            for node in tmp:\n",
    "                while e[node] and e[node][-1][1] <= curv:\n",
    "                    nxt = e[node][-1][0]\n",
    "                    uf.union(node,nxt)\n",
    "                    e[node].pop()\n",
    "            # 连通块内部点的个数\n",
    "            for node in tmp:\n",
    "                rt = uf.find(node)\n",
    "                vv.add(rt)\n",
    "                cc[rt] += 1\n",
    "            \n",
    "            for j in vv:\n",
    "                ans += cc[j] * (cc[j] - 1) // 2\n",
    "        return ans       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        maxcnt = [1] * n\n",
    "        fa = list(range(n))\n",
    "        def find(x: int)-> int:\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(x:int, y:int)-> int:\n",
    "            path = 0\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if vals[fx] > vals[fy]:\n",
    "                fa[fy] = fx\n",
    "            elif vals[fx] < vals[fy]:\n",
    "                fa[fx] = fy\n",
    "            else:\n",
    "                path = maxcnt[fx] * maxcnt[fy]\n",
    "                fa[fy] = fx\n",
    "                maxcnt[fx] += maxcnt[fy]\n",
    "            return path\n",
    "        ans = n\n",
    "        for x, y in sorted(edges, key = lambda p: max(vals[p[0]], vals[p[1]])):\n",
    "            ans += union(x, y)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    MAXN = 30001\n",
    "\n",
    "    def __init__(self):\n",
    "        self.father = [0] * self.MAXN\n",
    "        self.maxcnt = [1] * self.MAXN\n",
    "\n",
    "    def build(self, n):\n",
    "        # 初始化并查集\n",
    "        for i in range(n):\n",
    "            self.father[i] = i\n",
    "            self.maxcnt[i] = 1\n",
    "\n",
    "    def find(self, i):\n",
    "        # 寻找元素 i 所在集合的代表节点\n",
    "        if i != self.father[i]:\n",
    "            self.father[i] = self.find(self.father[i])\n",
    "        return self.father[i]\n",
    "\n",
    "    def union(self, x, y, vals):\n",
    "        # 合并两个集合\n",
    "        fx = self.find(x)\n",
    "        fy = self.find(y)\n",
    "        path = 0\n",
    "        if vals[fx] > vals[fy]:\n",
    "            self.father[fy] = fx\n",
    "        elif vals[fx] < vals[fy]:\n",
    "            self.father[fx] = fy\n",
    "        else:\n",
    "            # 两个集合的最大值相等\n",
    "            path = self.maxcnt[fx] * self.maxcnt[fy]\n",
    "            self.father[fy] = fx\n",
    "            self.maxcnt[fx] += self.maxcnt[fy]\n",
    "        return path\n",
    "\n",
    "    def numberOfGoodPaths(self, vals, edges):\n",
    "        n = len(vals)\n",
    "        self.build(n)\n",
    "        ans = n\n",
    "        # 核心排序：按照边上节点的最大值排序\n",
    "        edges.sort(key=lambda e: max(vals[e[0]], vals[e[1]]))\n",
    "        for edge in edges:\n",
    "            ans += self.union(edge[0], edge[1], vals)\n",
    "        return ans\n",
    "\n",
    "# 测试用例\n",
    "solution = Solution()\n",
    "\n",
    "# 例子1\n",
    "vals1 = [2, 1, 1, 2, 2, 1, 1, 2]\n",
    "edges1 = [\n",
    "    [0, 1],\n",
    "    [0, 2],\n",
    "    [1, 3],\n",
    "    [2, 4],\n",
    "    [2, 5],\n",
    "    [5, 6],\n",
    "    [6, 7]\n",
    "]\n",
    "print(solution.numberOfGoodPaths(vals1, edges1))\n",
    "\n",
    "# 例子2\n",
    "vals2 = [1, 2, 2, 3, 1, 2, 2, 1, 1, 3, 3, 3, 3]\n",
    "edges2 = [\n",
    "    [0, 1],\n",
    "    [0, 2],\n",
    "    [0, 3],\n",
    "    [1, 4],\n",
    "    [4, 7],\n",
    "    [4, 8],\n",
    "    [3, 5],\n",
    "    [3, 6],\n",
    "    [6, 9],\n",
    "    [6, 10],\n",
    "    [6, 11],\n",
    "    [9, 12]\n",
    "]\n",
    "print(solution.numberOfGoodPaths(vals2, edges2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        self.father = [0]*len(vals)\n",
    "        self.maxcnt = [0]*len(vals)\n",
    "        def build(n):\n",
    "            for i in range(n):\n",
    "                self.father[i] = i\n",
    "                self.maxcnt[i] = 1\n",
    "        def find(i):\n",
    "            if i!=self.father[i]:\n",
    "                self.father[i] = find(self.father[i])\n",
    "            return self.father[i]\n",
    "        def union(x,y,vals):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            path = 0\n",
    "            if vals[fx]>vals[fy]:\n",
    "                self.father[fy] = fx\n",
    "            elif vals[fy]>vals[fx]:\n",
    "                self.father[fx] = fy\n",
    "            else:\n",
    "                path += self.maxcnt[fx]*self.maxcnt[fy]\n",
    "                self.father[fy] = fx\n",
    "                self.maxcnt[fx] += self.maxcnt[fy]\n",
    "            return path\n",
    "        n = len(vals)\n",
    "        build(n)\n",
    "        ans = n\n",
    "        edges.sort(key=lambda a:max(vals[a[0]],vals[a[1]]))\n",
    "        for edge in edges:\n",
    "            ans += union(edge[0],edge[1],vals)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "         \n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        ans = n\n",
    "        UF = UnionFind(n)\n",
    "        for u in sorted(range(n), key=lambda i: vals[i]):\n",
    "            for v in g[u]:\n",
    "                ans += UF.unionSet(u, v, vals)\n",
    "        return ans \n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int, m: int=None) -> None:\n",
    "        self.parents = list(range(n))     \n",
    "        self.maxCount = [1] * n # maxCount[i]: number of vertices in the disjoint set rooted at vertex i with the same value as root (vals[i])        \n",
    "    \n",
    "    def findSet(self, i: int) -> int:\n",
    "        if self.parents[i] == i:\n",
    "            return i\n",
    "        else:\n",
    "            self.parents[i] = self.findSet(self.parents[i])\n",
    "            return self.parents[i]\n",
    "        \n",
    "    def unionSet(self, i: int, j: int, vals: List[int]) -> None:\n",
    "        \"\"\"若节点i和节点j来自2个不同的互斥集合，将两个集合合并 \n",
    "           T: amortized O(α(N))  effectively O(1)\n",
    "             first O(N) then O(1)\n",
    "        \"\"\"\n",
    "        goodPath = 0\n",
    "        ip, jp = self.findSet(i), self.findSet(j)\n",
    "\n",
    "        # 若两个节点来自同一互斥集合或节点j的根值大于节点i，返回\n",
    "        if ip == jp or vals[ip] < vals[jp]:\n",
    "            return goodPath \n",
    "\n",
    "        if vals[jp] == vals[ip]:\n",
    "            goodPath = self.maxCount[ip] * self.maxCount[jp]\n",
    "            self.maxCount[ip] += self.maxCount[jp]\n",
    "\n",
    "        # 节点值小的节点连到节点值大的树上\n",
    "        self.parents[jp] = ip\n",
    "\n",
    "        return goodPath\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        def find(u):\n",
    "            if fa[u] != u:\n",
    "                fa[u] = find(fa[u])\n",
    "            return fa[u]\n",
    "        n = len(vals)\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        ans = n\n",
    "        fa = list(range(n))                     # 以连通块中最大val值的节点作为代表元\n",
    "        size = [1] * n                          # size用于查询连通块中，与代表元pu同一val值的节点个数，注意代表元的val值是其连通块中最大val值\n",
    "        for u in sorted(range(n), key=lambda x: vals[x]):\n",
    "            for v in g[u]:\n",
    "                if vals[v] > vals[u]: continue  # 邻居节点值大的当前不进行合并\n",
    "                pu, pv = find(u), find(v)       # 找到这两个集合的代表元，注意pu节点的节点值一定是vals[u]，vals[u]是当前访问过所有节点的最大val\n",
    "                if pu == pv: continue           # 已经在一个连通块了，之前合并的时候计算过贡献了，不用再合并计算贡献了\n",
    "                if vals[pu] == vals[pv]:        # 只有两者连通块代表元的节点值一样大(满足条件1)，才能贡献好路径\n",
    "                    ans += size[pu] * size[pv]  # 乘法原理。因为pu连通块最大值节点与pv连通块最大值节点相连，路径上的节点值都小于等于最大值，满足条件2\n",
    "                    size[pu] += size[pv]        # 如果两者代表元的节点值一样大，将pv连通块最大值节点个数累加到pu连通块中；否则pu连通块个数不变\n",
    "                fa[pv] = pu                     # 由于是按照节点值从小到大遍历的，pu的节点值是vals[u]是当前最大的，所以pv合并到pu\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n=len(vals)\n",
    "        g=[[] for i in range(n)]\n",
    "        for i in edges:\n",
    "            g[i[0]].append(i[1])\n",
    "            g[i[1]].append(i[0])\n",
    "        \n",
    "        \n",
    "        parent=[i for i in range(n)]\n",
    "        \n",
    "        size=[1]*n\n",
    "        \n",
    "        def find(i):\n",
    "            if i==parent[i]:\n",
    "                return parent[i]\n",
    "            parent[i]=find(parent[i])\n",
    "            return parent[i]\n",
    "        \n",
    "        \n",
    "        ans=n\n",
    "        \n",
    "        for vx,x in sorted(zip(vals,range(n))):\n",
    "            rootx=find(x)\n",
    "            \n",
    "            for y in g[x]:\n",
    "                rooty=find(y)\n",
    "                \n",
    "                if rootx==rooty:\n",
    "                    continue\n",
    "                if vals[rooty]>vx:continue\n",
    "                if vals[rooty]==vx:\n",
    "                    ans+=size[rootx]*size[rooty]\n",
    "                    size[rootx]+=size[rooty]\n",
    "                parent[rooty]=rootx\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        p = [i for i in range(n)]\n",
    "        g = [[] for _ in range(n)]\n",
    "        size = [1] * n\n",
    "\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "\n",
    "        def find(x):\n",
    "            if x != p[x]:\n",
    "                p[x] = find(p[x])\n",
    "            return p[x]\n",
    "\n",
    "        res = n\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                fy = find(y)\n",
    "                if fx == fy or vals[fy] > vx: continue \n",
    "                if vals[fy] == vx:\n",
    "                    res += size[fx] * size[fy]\n",
    "                    size[fx] += size[fy]\n",
    "                p[fy] = fx \n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        p = [i for i in range(n)]\n",
    "        g = [[] for _ in range(n)]\n",
    "        size = [1] * n\n",
    "\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "\n",
    "        def find(x):\n",
    "            if x != p[x]:\n",
    "                p[x] = find(p[x])\n",
    "            return p[x]\n",
    "\n",
    "        res = n\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                fy = find(y)\n",
    "                if fx == fy or vals[fy] > vx: continue \n",
    "                if vals[fy] == vx:\n",
    "                    res += size[fx] * size[fy]\n",
    "                    size[fx] += size[fy]\n",
    "                p[fy] = fx \n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for s, e in edges:\n",
    "            g[e].append(s)\n",
    "            g[s].append(e)\n",
    "        # 列举每个节点作为出发点会超时======需要使用并查集，逐步创建连通集合\n",
    "        ans = n\n",
    "        fa = list(range(n))\n",
    "        size = [1 for _ in range(n)]\n",
    "        def find(x):\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                fy = find(y)\n",
    "                if fy == fx or vals[fy] > vx:\n",
    "                    continue\n",
    "                if vals[fy] == vx:\n",
    "                    ans += size[fx] * size[fy]\n",
    "                    size[fx] += size[fy]\n",
    "                fa[fy] = fx\n",
    "        return ans\n",
    "\n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, v: List[int], e: List[List[int]]) -> int:\n",
    "        n = len(v)\n",
    "        fa = [j for j in range(n)]\n",
    "        g = defaultdict(list)\n",
    "        for x, y in e:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x) \n",
    "        size = [1]*n \n",
    "        def find(x):     \n",
    "            if fa[x]!=x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x] \n",
    "        self.ans = n  \n",
    "        for j in sorted(range(n), key=lambda x:v[x]):\n",
    "            for k in g[j]:\n",
    "                if v[j]<v[k]:\n",
    "                    pass \n",
    "                else: \n",
    "                    fa[k] = find(k) \n",
    "                    fa[j] = find(j)\n",
    "                    if v[fa[k]] < v[fa[j]]: \n",
    "                        fa[fa[k]] = fa[j] \n",
    "                    else:\n",
    "                        if fa[j] == fa[k]:\n",
    "                            pass \n",
    "                        else:\n",
    "                            self.ans += size[fa[k]]*size[fa[j]]\n",
    "                            size[fa[j]] += size[fa[k]]\n",
    "                            fa[fa[k]] = fa[j] \n",
    "        return self.ans \n",
    "                            \n",
    "                            \n",
    "\n",
    "          \n",
    "                    \n",
    " \n",
    "                                \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        pa = list(range(n))\n",
    "        we = [1]*n\n",
    "        def find(x):\n",
    "            if pa[x] != x:\n",
    "                pa[x] = find(pa[x])\n",
    "            return pa[x]\n",
    "        \n",
    "        ans = n\n",
    "        for xv,x in sorted(zip(vals,range(n))):\n",
    "            xp = find(x)\n",
    "            for y in g[x]:\n",
    "                if vals[y] > xv: continue\n",
    "                yp = find(y)\n",
    "                if yp == xp: continue\n",
    "                pa[yp] = xp\n",
    "                if vals[yp] == xv:\n",
    "                    ans += we[xp] * we[yp]\n",
    "                    we[xp] += we[yp]\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def merge(from_: int, to: int) -> None:\n",
    "            from_ = find(from_)\n",
    "            to = find(to)\n",
    "            if from_ != to:\n",
    "                fa[from_] = to\n",
    "                size[to] += size[from_] \n",
    "        \n",
    "        ans = n\n",
    "        for val_x, x in sorted(zip(vals, range(n))):\n",
    "            px = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == px or val_x < vals[y]: continue\n",
    "                if val_x == vals[y]:\n",
    "                    ans += size[y] * size[px]\n",
    "                    size[px] += size[y]\n",
    "                fa[y] = px\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        #copy 0x3\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        fa = list(range(n))\n",
    "        size = [1]*n\n",
    "\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        ans  = n\n",
    "        for vx,x in sorted(zip(vals,range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y ==fx or vals[y]>vx:\n",
    "                    continue\n",
    "                if vals[y] == vx:\n",
    "                    ans += size[fx]*size[y]\n",
    "                    size[fx] += size[y]\n",
    "                fa[y] = fx\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n=len(vals)\n",
    "        g=[[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        fa=list(range(n))\n",
    "        size=[1]*n\n",
    "        def find(x):\n",
    "            if fa[x]!=x:\n",
    "                fa[x]=find(fa[x])\n",
    "            return fa[x]\n",
    "        ans=n\n",
    "        for vx,x in sorted(zip(vals,range(n))):\n",
    "            fx=find(x)\n",
    "            for y in g[x]:\n",
    "                y=find(y)\n",
    "                if y==fx or vals[y]>vx:\n",
    "                    continue\n",
    "                if vals[y]==vx:\n",
    "                    ans+=size[fx]*size[y]\n",
    "                    size[fx]+=size[y]\n",
    "                fa[y]=fx\n",
    "        return ans                            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "#好路径的数目，并查集，输入：vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]]输出：6\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x) \n",
    "        # 并查集\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:continue  #只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals, edges):\n",
    "        # 建图\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                # y==fx说明会构成环\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "       \n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        size = [1] * n # size[x]表示节点值等于vals[x]的节点个数， 如果按照节点值从小到大合并，size[x]也是连通块内等于最大节点值的节点个数\n",
    "        def find(x:int)-> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue # 只考虑最大节点值不炒作vx的连通块\n",
    "                if vals[y] == vx: # 可以构成好路径\n",
    "                    ans += size[fx] * size[y] # 乘法原理\n",
    "                    size[fx] += size[y] # 统计连通块内节点值等于vx的节点个数\n",
    "                fa[y] = fx # 把小的节点值合并到大的节点值上\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        fa = list(range(n))\n",
    "        sz = [1] * n\n",
    "\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        # 单个节点也视为一条合法的路径\n",
    "        ans = n\n",
    "        for val_x, x in sorted(zip(vals, range(n))):\n",
    "            # 由于是从小到大排列，val_x是当前可以合并到连通块中，的最大节点值\n",
    "            fx = find(x) # x所处集合（连通块）的代表元\n",
    "            for y in g[x]:\n",
    "                fy = find(y)\n",
    "                if fy==fx or vals[fy] > val_x:\n",
    "                    continue\n",
    "                # 如果恰好相等，就找到了长度大于1的好路径\n",
    "                if vals[fy] == val_x:\n",
    "                    ans += sz[fy] * sz[fx]\n",
    "                    sz[fx] += sz[fy]\n",
    "                \n",
    "                fa[fy] = fx # 合并，从小的合并到大的\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def check(self, a, b):\n",
    "        \"\"\"\n",
    "        check if a good path starting from a, ending at b exists, by dfs.\n",
    "        @a, @b is node index where a!=b, and vals[a] == vals[b].\n",
    "        \"\"\"\n",
    "        stack = [a]\n",
    "        visited = [False] * self.N\n",
    "        while stack:\n",
    "            # print(stack)\n",
    "            # print(visited)\n",
    "            cur = stack.pop()\n",
    "            if cur == b:\n",
    "                return True\n",
    "            if visited[cur]:\n",
    "                continue\n",
    "            visited[cur] = True\n",
    "            if self.vals[cur] > self.vals[a]:\n",
    "                continue\n",
    "            for neibour in self.adj[cur]:\n",
    "                if not visited[neibour]:\n",
    "                    stack.append(neibour)\n",
    "        return False\n",
    "        \n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        # self.N = len(vals)\n",
    "        # self.vals = vals\n",
    "        # val_dict = {}\n",
    "        # for i in range(self.N):\n",
    "        #     if vals[i] in val_dict:\n",
    "        #         val_dict[vals[i]].append(i)\n",
    "        #     else:\n",
    "        #         val_dict[vals[i]] = [i]\n",
    "        # self.adj = [[] for _ in range(self.N)]\n",
    "        # for a, b in edges:\n",
    "        #     self.adj[a].append(b)\n",
    "        #     self.adj[b].append(a)\n",
    "        # ans = self.N\n",
    "        # for val in val_dict:\n",
    "        #     equaled_nodes = val_dict[val]\n",
    "        #     num_equaled_nodes = len(equaled_nodes)\n",
    "        #     if num_equaled_nodes >= 2:\n",
    "        #         for i in range(num_equaled_nodes):\n",
    "        #             for j in range(i + 1, num_equaled_nodes):\n",
    "        #                 # print(equaled_nodes[i], equaled_nodes[j])\n",
    "        #                 if self.check(equaled_nodes[i], equaled_nodes[j]):\n",
    "        #                     ans += 1\n",
    "        # return ans\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        N = len(vals)\n",
    "        adj = [[] for _ in range(N)]\n",
    "        for a, b in edges:\n",
    "            adj[a].append(b)\n",
    "            adj[b].append(a)\n",
    "        fa = list(range(N))\n",
    "        size = [1] * N\n",
    "\n",
    "        ans = N\n",
    "        for vx, x in sorted(zip(vals, range(N))):  # sort [vx, x] pair \n",
    "            fx = find(x)\n",
    "            for y in adj[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    # unioned or larger neibour, \n",
    "                    # for unioned, good path had been counted.\n",
    "                    # for larger neibour, union later.\n",
    "                    continue\n",
    "                if vals[y] == vx:\n",
    "                    # found a food path, that start/end at y/x, where x != x and all intermediate node's < vx.\n",
    "                    ans += size[fx] * size[y]\n",
    "                    size[fx] += size[y]\n",
    "                fa[y] = fx  # union\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        fa = list(range(n))\n",
    "        size = [1] * n\n",
    "        def find(x:int)->int:\n",
    "            if fa[x]!=x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        ans = n\n",
    "        for vx,x in sorted(zip(vals,range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue\n",
    "                if vals[y] == vx:\n",
    "                    ans += size[fx]*size[y]\n",
    "                    size[fx]+= size[y]\n",
    "                fa[y]=fx\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                fy = find(y)\n",
    "                if fy == fx or vals[fy] > vals[fx]:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[fy] == vals[fx]:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[fy]  # 乘法原理\n",
    "                    # 统计连通块内节点值等于 vx 的节点个数\n",
    "                    size[fx] += size[fy] \n",
    "                fa[fy] = fx  # 把小的节点值合并到大的节点值上\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)  # 建图\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\n",
    "        size = [1] * n\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n  # 单个节点的好路径\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            fx = find(x)\n",
    "            for y in g[x]:\n",
    "                y = find(y)\n",
    "                if y == fx or vals[y] > vx:\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\n",
    "                if vals[y] == vx:  # 可以构成好路径\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\n",
    "        return ans\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\r\n",
    "        n = len(vals)\r\n",
    "        g = [[] for _ in range(n)]\r\n",
    "        for x, y in edges:\r\n",
    "            g[x].append(y)\r\n",
    "            g[y].append(x)  # 建图\r\n",
    "\r\n",
    "        # 并查集模板\r\n",
    "        fa = list(range(n))\r\n",
    "        # size[x] 表示节点值等于 vals[x] 的节点个数，\r\n",
    "        # 如果按照节点值从小到大合并，size[x] 也是连通块内的等于最大节点值的节点个数\r\n",
    "        size = [1] * n\r\n",
    "        def find(x: int) -> int:\r\n",
    "            if fa[x] != x:\r\n",
    "                fa[x] = find(fa[x])\r\n",
    "            return fa[x]\r\n",
    "\r\n",
    "        ans = n  # 单个节点的好路径\r\n",
    "        for vx, x in sorted(zip(vals, range(n))):\r\n",
    "            fx = find(x)\r\n",
    "            for y in g[x]:\r\n",
    "                y = find(y)\r\n",
    "                if y == fx or vals[y] > vx:\r\n",
    "                    continue  # 只考虑最大节点值不超过 vx 的连通块\r\n",
    "                if vals[y] == vx:  # 可以构成好路径\r\n",
    "                    ans += size[fx] * size[y]  # 乘法原理\r\n",
    "                    size[fx] += size[y]  # 统计连通块内节点值等于 vx 的节点个数\r\n",
    "                fa[y] = fx  # 把小的节点值合并到大的节点值上\r\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "'''\n",
    "输出好路径的数目\n",
    "开始和结束节点的值相同，并且是最大的\n",
    "按照节点的值从小到大枚举，慢慢合并\n",
    "如果同一个值多次加入到了一个集合中，那么，直接计算组合数，代表元素随意选取一个最大的值即可\n",
    "'''\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n=len(vals)\n",
    "        ans=n\n",
    "        g=[[] for _ in range(n)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        fa=list(range(n))\n",
    "        sz=[1]*n\n",
    "        mx=[x for x in vals]\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        for i,x in sorted(enumerate(vals),key=lambda z:z[1]):\n",
    "            #使用一边遍历的思想吧\n",
    "            for y in g[i]:\n",
    "                j=find(y)\n",
    "                fx=find(i)\n",
    "                if fx==j:continue\n",
    "                if mx[j]>x:continue\n",
    "                '''\n",
    "                这样的写法是不对的，可能有多个连通块\n",
    "                '''\n",
    "                if mx[j]==x:\n",
    "                    \n",
    "                    ans+=sz[fx]*sz[j]\n",
    "                    sz[fx]+=sz[j]\n",
    "                fa[j]=fx\n",
    "                mx[j]=x\n",
    "                        \n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        n = len(vals)\n",
    "        fa = list(range(n))\n",
    "\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        size = [1]*n\n",
    "        ans = n\n",
    "        for x,u in sorted(zip(vals,range(n))):\n",
    "            fu = find(u) # u所在集合\n",
    "            for v in g[u]:\n",
    "                fv = find(v)\n",
    "                if fv == fu or vals[fv] > x:\n",
    "                    continue # 如果本来就连着，跳过；如果v所在连通块的最大值更大，没法合并，跳过\n",
    "                if vals[fv] == x:\n",
    "                    ans += size[fv]*size[fu]\n",
    "                    # x是此时最大，从小的合并到大的\n",
    "                    size[fu] += size[fv]\n",
    "                fa[fv] = fu\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        fa = list(range(n))\n",
    "        g = defaultdict(list)\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        size = [1]*n\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        ans = n\n",
    "        for x,u in sorted(zip(vals,range(n))):\n",
    "            fu = find(u)\n",
    "            for v in g[u]:\n",
    "                fv = find(v)\n",
    "                if fv == fu or vals[fv] > x:\n",
    "                    continue\n",
    "                if vals[fv] == vals[fu]:# vals[fu] == x\n",
    "                    ans += size[fv]*size[fu]\n",
    "                    size[fu] += size[fv]\n",
    "                fa[fv] = fu\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Dsu:\n",
    "    def __init__(self, n):\n",
    "        self.pa = list(range(n))\n",
    "        self.size = [1] * n\n",
    "\n",
    "    def find(self, x):\n",
    "        if x != self.pa[x]:\n",
    "            self.pa[x] = self.find(self.pa[x])\n",
    "        return self.pa[x]\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        dsu = Dsu(n)\n",
    "        g = defaultdict(list)\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        ans = n\n",
    "        for vx, x in sorted(zip(vals, range(n))):\n",
    "            px = dsu.find(x)\n",
    "            for y in g[x]:\n",
    "                py = dsu.find(y)\n",
    "                if py == px or vals[py] > vx:\n",
    "                    continue\n",
    "                if vals[py] == vx:\n",
    "                    ans += dsu.size[px] * dsu.size[py]\n",
    "                    dsu.size[px] += dsu.size[py]\n",
    "                dsu.pa[py] = px\n",
    "                    \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        n = len(vals)\n",
    "        fa = list(range(n))\n",
    "\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        \n",
    "        size = [1]*n\n",
    "        ans = n\n",
    "        for x,u in sorted(zip(vals,range(n))):\n",
    "            fu = find(u) # u所在集合\n",
    "            for v in g[u]:\n",
    "                fv = find(v)\n",
    "                if fv == fu or vals[fv] > x:\n",
    "                    continue # 如果本来就连着，跳过；如果v所在连通块的最大值更大，没法合并，跳过\n",
    "                if vals[fv] == x:\n",
    "                    ans += size[fv]*size[fu]\n",
    "                    # x是此时最大，从小的合并到大的\n",
    "                    size[fu] += size[fv]\n",
    "                    fa[fv] = fu\n",
    "                fa[fv] = fu\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        self.father = [0]*30001\n",
    "        self.maxcnt = [0]*30001\n",
    "        def build(n):\n",
    "            for i in range(n):\n",
    "                self.father[i] = i\n",
    "                self.maxcnt[i] = 1\n",
    "        def find(i):\n",
    "            if i!=self.father[i]:\n",
    "                self.father[i] = find(self.father[i])\n",
    "            return self.father[i]\n",
    "        def union(x,y,vals):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            path = 0\n",
    "            if vals[fx]>vals[fy]:\n",
    "                self.father[fy] = fx\n",
    "            elif vals[fy]>vals[fx]:\n",
    "                self.father[fx] = fy\n",
    "            else:\n",
    "                path += self.maxcnt[fx]*self.maxcnt[fy]\n",
    "                self.father[fy] = fx\n",
    "                self.maxcnt[fx] += self.maxcnt[fy]\n",
    "            return path\n",
    "        n = len(vals)\n",
    "        build(n)\n",
    "        ans = n\n",
    "        edges.sort(key=lambda a:max(vals[a[0]],vals[a[1]]))\n",
    "        for edge in edges:\n",
    "            ans += union(edge[0],edge[1],vals)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        maxcnt = [1] * n\n",
    "        fa = list(range(n))\n",
    "        def find(x: int)-> int:\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(x:int, y:int)-> int:\n",
    "            path = 0\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if vals[fx] > vals[fy]:\n",
    "                fa[fy] = fx\n",
    "            elif vals[fx] < vals[fy]:\n",
    "                fa[fx] = fy\n",
    "            else:\n",
    "                path = maxcnt[fx] * maxcnt[fy]\n",
    "                fa[fy] = fx\n",
    "                maxcnt[fx] += maxcnt[fy]\n",
    "            return path\n",
    "        ans = n\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[max(vals[u], vals[v])].append((u, v))\n",
    "        for val, edge in sorted(g.items()):\n",
    "            for x, y in edge:\n",
    "                ans += union(x, y)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "# -*- coding: utf-8 -*-\n",
    "from typing import List, Tuple, Optional\n",
    "from collections import deque, Counter\n",
    "from queue import PriorityQueue\n",
    "import math\n",
    "from functools import lru_cache\n",
    "from sortedcontainers import SortedDict, SortedSet\n",
    "import random\n",
    "import copy\n",
    "\n",
    "import sys\n",
    "sys.setrecursionlimit(9999999)\n",
    "\n",
    "MOD = 10**9 + 7\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "\n",
    "        link = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            link[a].append(b)\n",
    "            link[b].append(a)\n",
    "\n",
    "        mm = SortedDict()\n",
    "        for x in range(n):\n",
    "            if vals[x] not in mm:\n",
    "                mm[vals[x]] = []\n",
    "            mm[vals[x]].append(x)\n",
    "\n",
    "        p = [i for i in range(n)]\n",
    "\n",
    "        def get_rt(x):\n",
    "            if p[x] != x:\n",
    "                p[x] = get_rt(p[x])\n",
    "            return p[x]\n",
    "\n",
    "        def merge(a, b):\n",
    "            aa, bb = get_rt(a), get_rt(b)\n",
    "            if aa != bb:\n",
    "                p[aa] = bb\n",
    "\n",
    "        ans = n  # 初始值是只有一个点构成的路径，总共n条\n",
    "        for key, nodes in mm.items():\n",
    "            for node in nodes:\n",
    "                for ne in link[node]:\n",
    "                    if vals[ne] <= key:\n",
    "                        merge(ne, node)\n",
    "\n",
    "            c = Counter()\n",
    "            for node in nodes:\n",
    "                c[get_rt(node)] += 1\n",
    "            for v in c.values():\n",
    "                ans += v * (v - 1) // 2\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        def pa(a):\n",
    "            if arr[a] != a:\n",
    "                arr[a] = pa(arr[a])\n",
    "            return arr[a]\n",
    "        \n",
    "        def union(a, b):\n",
    "            arr[pa(a)] = pa(b)\n",
    "\n",
    "        n = len(vals)\n",
    "        arr = list(range(n))\n",
    "        g = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        ans = 0\n",
    "        cnt = defaultdict(lambda :1)\n",
    "        for u, a in sorted(enumerate(vals), key = lambda x: x[1]):\n",
    "            pu = pa(u)\n",
    "            cur = cnt[pu]\n",
    "            for v in g[u]:                \n",
    "                pv = pa(v)\n",
    "                if a < vals[pv]:\n",
    "                    continue\n",
    "                if a == vals[pv] and pv != pu:                    \n",
    "                    ans += cur * cnt[pv]\n",
    "                    cur += cnt[pv]\n",
    "                union(pv, pu)\n",
    "            cnt[pu] = cur            \n",
    "        return ans + n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(vals)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        pa = list(range(n))\n",
    "        we = [1] * n\n",
    "        \n",
    "        def find(x):\n",
    "            if pa[x] != x:\n",
    "                pa[x] = find(pa[x])\n",
    "            return pa[x]\n",
    "        \n",
    "        locs = defaultdict(list)\n",
    "        for i,x in enumerate(vals): locs[x].append(i)\n",
    "        \n",
    "        ans = 0\n",
    "        for xv,l in sorted(locs.items()):\n",
    "            for x in l:\n",
    "                xr = find(x)\n",
    "                for y in g[x]:\n",
    "                    if vals[y] > xv: continue\n",
    "                    yr = find(y)\n",
    "                    if xr == yr: continue\n",
    "                    pa[yr] = xr\n",
    "                    if vals[yr] == xv:\n",
    "                        ans += we[xr] * we[yr]\n",
    "                        we[xr] += we[yr]\n",
    "        return ans + n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], arr: List[List[int]]) -> int:\n",
    "        edge=defaultdict(list)\n",
    "        node=defaultdict(list)\n",
    "        for i,v in enumerate(vals):\n",
    "            node[v].append(i)\n",
    "        for a,b in arr:\n",
    "            edge[max(vals[a],vals[b])].append([a,b])\n",
    "        res=0\n",
    "        fa=list(range(len(vals)))\n",
    "        def find(a):\n",
    "            if fa[a]!=a:\n",
    "                fa[a]=find(fa[a])\n",
    "            return fa[a]\n",
    "        for v in sorted(node.keys()):\n",
    "            for a,b in edge[v]:\n",
    "                fa[find(a)]=find(b)\n",
    "            cnt=defaultdict(int)\n",
    "            for a in node[v]:\n",
    "                cnt[find(a)]+=1\n",
    "                res+=cnt[find(a)]\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "\n",
    "        f=defaultdict(list)\n",
    "        for i,v in enumerate(vals):\n",
    "            f[v].append(i)\n",
    "        a=[(v,arr) for v,arr in f.items()]\n",
    "        a.sort(key=lambda x:x[0])\n",
    "\n",
    "        g=defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        cnt=Counter()\n",
    "\n",
    "        n=len(vals)\n",
    "        parent=list(range(n))\n",
    "        def find(x):\n",
    "            if x!=parent[x]:\n",
    "                parent[x]=find(parent[x])\n",
    "            return parent[x]\n",
    "        def union(x,y):\n",
    "            a,b=find(x),find(y)\n",
    "            if a!=b:parent[b]=a\n",
    "\n",
    "        res=0\n",
    "        for v,arr in a:\n",
    "            for i in arr:\n",
    "                for j in g[i]:\n",
    "                    if vals[j]<=vals[i]:\n",
    "                        union(i,j)\n",
    "            cnt=Counter()\n",
    "            for i in arr:\n",
    "                root=find(i)\n",
    "                res+=cnt[root]\n",
    "                cnt[root]+=1\n",
    "        \n",
    "        return res+n\n",
    "\n",
    "            \n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/problems/number-of-good-paths/\n",
    "# 一条 好路径需要满足以下条件：\n",
    "\n",
    "# !开始节点和结束节点的值 相同。\n",
    "# !开始节点和结束节点中间的所有节点值都小于等于开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。\n",
    "# 请你返回不同好路径的数目。\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[int]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n=len(vals)\n",
    "        g=[[] for _ in range(n)]\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        tree,root=centroidDecomposition(n,g)\n",
    "        self.res=0\n",
    "        def dfs1(u): #点分治\n",
    "            use[u]=True\n",
    "            for v in tree[u]:\n",
    "                dfs1(v)\n",
    "            use[u]=False\n",
    "            self.cnt=defaultdict(int)\n",
    "            self.cnt[vals[u]]+=1\n",
    "            for v in g[u]:\n",
    "                if not use[v]:\n",
    "                    dfs2(v,u,vals[u])\n",
    "                    dfs3(v,u,vals[u])\n",
    "            #print(u,ban,g[u],self.cnt,self.res)\n",
    "        \n",
    "        def dfs2(u,fa,maxval):  #累加答案\n",
    "            if vals[u]>=maxval:\n",
    "                self.res+=self.cnt[vals[u]]\n",
    "                maxval=vals[u]\n",
    "            for v in g[u]:\n",
    "                if v!=fa and not use[v]:\n",
    "                    dfs2(v,u,maxval)\n",
    "\n",
    "        def dfs3(u,fa,maxval):  #累加路径到哈希表里\n",
    "            if vals[u]>=maxval:\n",
    "                self.cnt[vals[u]]+=1\n",
    "                maxval=vals[u]\n",
    "            for v in g[u]:\n",
    "                if v!=fa and not use[v]:\n",
    "                    dfs3(v,u,maxval)        \n",
    "        use=[False]*n\n",
    "        dfs1(root)\n",
    "        return self.res+n\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/problems/number-of-good-paths/\n",
    "# 一条 好路径需要满足以下条件：\n",
    "\n",
    "# !开始节点和结束节点的值 相同。\n",
    "# !开始节点和结束节点中间的所有节点值都小于等于开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。\n",
    "# 请你返回不同好路径的数目。\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[int]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n=len(vals)\n",
    "        g=[[] for _ in range(n)]\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        tree,root=centroidDecomposition(n,g)\n",
    "        self.res=0\n",
    "        def dfs1(u): \n",
    "            use[u]=True\n",
    "            for v in tree[u]:\n",
    "                dfs1(v)\n",
    "            use[u]=False\n",
    "            self.cnt=defaultdict(int)\n",
    "            self.cnt[vals[u]]+=1\n",
    "            for v in g[u]:\n",
    "                if not use[v]:\n",
    "                    dfs2(v,u,vals[u],True)\n",
    "                    dfs2(v,u,vals[u],False)\n",
    "        def dfs2(u,fa,maxval,flag):  \n",
    "            if vals[u]>=maxval:\n",
    "                if flag:\n",
    "                    self.res+=self.cnt[vals[u]]\n",
    "                else:\n",
    "                    self.cnt[vals[u]]+=1\n",
    "                maxval=vals[u]\n",
    "            for v in g[u]:\n",
    "                if v!=fa and not use[v]:\n",
    "                    dfs2(v,u,maxval,flag)     \n",
    "        use=[False]*n\n",
    "        dfs1(root)\n",
    "        return self.res+n\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/problems/number-of-good-paths/\n",
    "# 一条 好路径需要满足以下条件：\n",
    "\n",
    "# !开始节点和结束节点的值 相同。\n",
    "# !开始节点和结束节点中间的所有节点值都小于等于开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。\n",
    "# 请你返回不同好路径的数目。\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[int]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        def collect(cur: int, pre: int, sub: \"DefaultDict[int, int]\", maxValue: int) -> None:\n",
    "            \"\"\"统计子树内的答案.\"\"\"\n",
    "            curValue = vals[cur]\n",
    "            if curValue >= maxValue:\n",
    "                sub[curValue] += 1\n",
    "                maxValue = curValue\n",
    "            for next in tree[cur]:\n",
    "                if next != pre and not removed[next]:\n",
    "                    collect(next, cur, sub, maxValue)\n",
    "        def decomposition(cur: int, pre: int) -> None:\n",
    "            nonlocal res\n",
    "            removed[cur] = True\n",
    "            for next in centTree[cur]:  # 点分树的子树中的答案(不经过重心)\n",
    "                if not removed[next]:\n",
    "                    decomposition(next, cur)\n",
    "            removed[cur] = False\n",
    "            counter = defaultdict(int, {vals[cur]: 1})  # 经过重心的路径\n",
    "            for next in tree[cur]:\n",
    "                if next == pre or removed[next]:\n",
    "                    continue\n",
    "                sub = defaultdict(int)\n",
    "                collect(next, cur, sub, vals[cur])\n",
    "                for v, count in sub.items():\n",
    "                    res += count * counter[v]\n",
    "                    counter[v] += count\n",
    "        n = len(edges) + 1\n",
    "        tree = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            tree[u].append(v)\n",
    "            tree[v].append(u)\n",
    "        centTree, root = centroidDecomposition(n, tree)\n",
    "        removed = [False] * n\n",
    "        res = 0\n",
    "        decomposition(root, -1)\n",
    "        return res + n\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/problems/number-of-good-paths/\n",
    "# 一条 好路径需要满足以下条件：\n",
    "\n",
    "# !开始节点和结束节点的值 相同。\n",
    "# !开始节点和结束节点中间的所有节点值都小于等于开始节点的值（也就是说开始节点的值应该是路径上所有节点的最大值）。\n",
    "# 请你返回不同好路径的数目。\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[int]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        def collect(cur: int, pre: int, sub: \"DefaultDict[int, int]\", maxValue: int) -> None:\n",
    "            \"\"\"统计子树内的答案.\"\"\"\n",
    "            curValue = vals[cur]\n",
    "            if curValue >= maxValue:\n",
    "                sub[curValue] += 1\n",
    "                maxValue = curValue\n",
    "            for next in tree[cur]:\n",
    "                if next != pre and not removed[next]:\n",
    "                    collect(next, cur, sub, maxValue)\n",
    "        def decomposition(cur: int, pre: int) -> None:\n",
    "            nonlocal res\n",
    "            removed[cur] = True\n",
    "            for next in centTree[cur]:  # 点分树的子树中的答案(不经过重心)\n",
    "                if not removed[next]:\n",
    "                    decomposition(next, cur)\n",
    "            removed[cur] = False\n",
    "            counter = defaultdict(int, {vals[cur]: 1})  # 经过重心的路径\n",
    "            for next in tree[cur]:\n",
    "                if next == pre or removed[next]:\n",
    "                    continue\n",
    "                sub = defaultdict(int)\n",
    "                collect(next, cur, sub, vals[cur])\n",
    "                for v, count in sub.items():\n",
    "                    res += count * counter[v]\n",
    "                    counter[v] += count\n",
    "        n = len(vals)\n",
    "        tree = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            tree[u].append(v)\n",
    "            tree[v].append(u)\n",
    "        centTree, root = centroidDecomposition(n, tree)\n",
    "        removed = [False] * n\n",
    "        res = 0\n",
    "        decomposition(root, -1)\n",
    "        return res + n\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def numberOfGoodPaths(self, vals: List[int], edges: List[List[int]]) -> int:\n",
    "        n = len(edges) + 1\n",
    "        tree = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            tree[u].append(v)\n",
    "            tree[v].append(u)\n",
    "\n",
    "        centTree, root = centroidDecomposition(n, tree)\n",
    "        removed = [False] * n\n",
    "        res = 0\n",
    "\n",
    "        def collect(cur: int, pre: int, sub: \"DefaultDict[int, int]\", maxValue: int) -> None:\n",
    "            \"\"\"统计子树内的答案.\"\"\"\n",
    "            curValue = vals[cur]\n",
    "            if curValue >= maxValue:\n",
    "                sub[curValue] += 1\n",
    "                maxValue = curValue\n",
    "            for next in tree[cur]:\n",
    "                if next != pre and not removed[next]:\n",
    "                    collect(next, cur, sub, maxValue)\n",
    "\n",
    "        def decomposition(cur: int, pre: int) -> None:\n",
    "            nonlocal res\n",
    "            removed[cur] = True\n",
    "            for next in centTree[cur]:  # 点分树的子树中的答案(不经过重心)\n",
    "                if not removed[next]:\n",
    "                    decomposition(next, cur)\n",
    "            removed[cur] = False\n",
    "\n",
    "            counter = defaultdict(int, {vals[cur]: 1})  # 经过重心的路径\n",
    "            for next in tree[cur]:\n",
    "                if next == pre or removed[next]:\n",
    "                    continue\n",
    "                sub = defaultdict(int)\n",
    "                collect(next, cur, sub, vals[cur])\n",
    "                for v, count in sub.items():\n",
    "                    res += count * counter[v]\n",
    "                    counter[v] += count\n",
    "\n",
    "        decomposition(root, -1)\n",
    "        return res + n\n",
    "\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[int]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
